#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
from std_msgs.msg import String
import cv2
import numpy as np
from cv_bridge import CvBridge
import mediapipe as mp

class GestureRecognizer(Node):
    def __init__(self):
        super().__init__('gesture_recognizer')
        
        # 初始化MediaPipe手部模型
        self.mp_hands = mp.solutions.hands
        self.hands = self.mp_hands.Hands(
            static_image_mode=False,
            max_num_hands=2,
            min_detection_confidence=0.7,
            min_tracking_confidence=0.5)
        
        # 初始化ROS2组件
        self.bridge = CvBridge()
        self.subscription = self.create_subscription(
            CompressedImage,
            '/camera/image_raw',
            self.image_callback,
            10)
        self.publisher = self.create_publisher(String, '/gesture_recognition', 10)
        
        self.get_logger().info("手势识别节点已启动，等待图像数据...")

    def count_fingers(self, hand_landmarks):
        """识别手势并返回结果"""
        tip_ids = [4, 8, 12, 16, 20]  # 拇指, 食指, 中指, 无名指, 小指
        
        fingers = []
        
        # 拇指判断 (比较x坐标)
        if hand_landmarks.landmark[tip_ids[0]].x < hand_landmarks.landmark[tip_ids[0]-1].x:
            fingers.append(1)
        else:
            fingers.append(0)
        
        # 其他四指判断 (比较y坐标)
        for id in range(1, 5):
            if hand_landmarks.landmark[tip_ids[id]].y < hand_landmarks.landmark[tip_ids[id]-2].y:
                fingers.append(1)
            else:
                fingers.append(0)
        
        total_fingers = fingers.count(1)
        
        # 特殊手势判断
        if fingers == [0, 1, 0, 0, 0]:
            return "1"
        elif fingers == [0, 1, 1, 0, 0]:
            return "2"
        elif fingers == [0, 1, 1, 1, 0]:
            return "3"
        elif fingers == [0, 1, 1, 1, 1]:
            return "4"
        elif fingers == [1, 1, 1, 1, 1]:
            return "5"
        elif fingers == [0, 0, 0, 0, 0]:
            return "0"
        else:
            return str(total_fingers)

    def image_callback(self, msg):
        """处理图像回调"""
        try:
            # 将压缩图像转换为OpenCV格式
            cv_image = self.bridge.compressed_imgmsg_to_cv2(msg, "bgr8")
        except Exception as e:
            self.get_logger().error(f"图像转换错误: {e}")
            return
        
        # 转换颜色空间并处理
        image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
        results = self.hands.process(image)
        
        # 检测到手部时
        if results.multi_hand_landmarks:
            for hand_landmarks in results.multi_hand_landmarks:
                # 获取手势识别结果
                gesture = self.count_fingers(hand_landmarks)
                
                # 发布识别结果
                gesture_msg = String()
                gesture_msg.data = gesture
                self.publisher.publish(gesture_msg)
                
                self.get_logger().info(f"检测到手势: {gesture}", throttle_duration_sec=1)
                
                # 在图像上绘制关键点（可选，调试用）
                mp.solutions.drawing_utils.draw_landmarks(
                    cv_image, hand_landmarks, self.mp_hands.HAND_CONNECTIONS)
                
                # 显示识别结果（可选，调试用）
                cv2.putText(cv_image, f"Gesture: {gesture}", (10, 30),
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 显示图像（可选，调试用）
        cv2.imshow('Gesture Recognition', cv_image)
        cv2.waitKey(1)

    def __del__(self):
        """析构函数"""
        self.hands.close()
        cv2.destroyAllWindows()

def main(args=None):
    rclpy.init(args=args)
    gesture_recognizer = GestureRecognizer()
    rclpy.spin(gesture_recognizer)
    gesture_recognizer.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()